import os
from dotenv import load_dotenv
from langchain.chat_models.glm import ChatGLM
from nodes import NovelNodes
from workflow import NovelWorkflow

load_dotenv()

class Config:
    GLM_API_KEY = os.getenv('GLM_API_KEY')
    GLM_API_BASE = os.getenv('GLM_API_BASE', 'https://open.bigmodel.cn/api/paas/v3/')

def main():
    # 初始化GLM-4模型
    llm = ChatGLM(
        api_key=Config.GLM_API_KEY,
        base_url=Config.GLM_API_BASE,
        model_name="glm-4"
    )

    # 构建工作流组件
    nodes = NovelNodes(llm)
    workflow = NovelWorkflow(nodes)
    app = workflow.compile_workflow()

    # 初始化流程状态
    initial_state = {
        "outline": "",
        "iteration": 0,
        "current_chapter": "第一章"
    }

    # 执行工作流
    final_state = app.invoke(initial_state)

    # 章节循环生成逻辑
    for chapter_num in range(2, 6):
        chapter_state = {
            "outline": final_state["outline"],
            "current_chapter": f"第{chapter_num}章"
        }
        app.invoke(chapter_state)

if __name__ == "__main__":
    main()